"""
Description: A simple teleoperation example using Vive Ultimate tracker to control a robot in simulation.

Author: Jin Liu
Date: 2025/06/18
"""

from ezgl import tb, Matrix4x4
from ezrobot.sim import RobotClient
from ezrobot.utils import GraphRunner, HandEyeConfig
from ezrobot.motion import OnlineMotion
from ezrobot.sensor import ViveTracker, SensorManager

from pathlib import Path
BASE_DIR = Path(__file__).resolve().parent


if __name__ == '__main__':
    
    # vive tracker
    vive = SensorManager(ViveTracker(device_index=1), frequency=30, buffer_size=10, streaming=True)
    handeye_config = HandEyeConfig.load(BASE_DIR / "../configs/handeye_config.yaml")
   
    # sim robot and motion
    client = RobotClient()
    client.clear_items()
    client.toggle_traj(True)
    motion = OnlineMotion(client.get_pose(), vmax=0.4, amax=0.1)

    # 异步运动跟踪 OnlineMotion
    motion_runner = (
        GraphRunner()
        .add_node(0, motion.get_pose)
        .add_node(1, client.set_cartesian, -1)
        .run_async(frequency=50)
    )
    
    # 遥操作更新 OnlineMotion
    def get_target_pose(vive2hand: Matrix4x4 = Matrix4x4.fromTranslation(0.5, 0, 0)):
        """
        获取目标位姿, vive2hand: vive 到 ee 的变换
        """
        return handeye_config.base2map * vive.get_data() * vive2hand
    
    vive_teleop_runner = (
        GraphRunner()
        .add_node(0, get_target_pose)
        .add_node(1, motion.add_pose, -1)
        .add_node(2, lambda p: client.set_pose_axis(p * Matrix4x4.fromScale(4, 4, 4)), 0)
    )
    vive_teleop_runner.run(30)